#!/usr/bin/env python3
"""
数据库初始化脚本
"""

import os
import sys
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from alembic.config import Config
from alembic import command

def init_database():
    """初始化数据库"""
    print("🔧 初始化数据库...")
    
    # 确保数据目录存在
    data_dir = project_root / "data"
    data_dir.mkdir(exist_ok=True)
    print(f"📁 数据目录: {data_dir}")
    
    # 获取alembic配置
    alembic_cfg = Config(str(project_root / "alembic.ini"))
    
    try:
        # 创建迁移表
        print("📝 创建迁移表...")
        command.stamp(alembic_cfg, "head")
        
        print("✅ 数据库初始化成功！")
        return True
    except Exception as e:
        print(f"❌ 数据库初始化失败: {e}")
        return False

def create_migration(message="Initial migration"):
    """创建新的迁移"""
    print("🔧 创建数据库迁移...")
    
    project_root = Path(__file__).parent.parent
    alembic_cfg = Config(str(project_root / "alembic.ini"))
    
    try:
        command.revision(alembic_cfg, autogenerate=True, message=message)
        print("✅ 迁移创建成功！")
        return True
    except Exception as e:
        print(f"❌ 迁移创建失败: {e}")
        return False

def upgrade_database():
    """升级数据库"""
    print("🔧 升级数据库...")
    
    project_root = Path(__file__).parent.parent
    alembic_cfg = Config(str(project_root / "alembic.ini"))
    
    try:
        command.upgrade(alembic_cfg, "head")
        print("✅ 数据库升级成功！")
        return True
    except Exception as e:
        print(f"❌ 数据库升级失败: {e}")
        return False

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="数据库初始化工具")
    parser.add_argument("--create-migration", action="store_true", help="创建新的迁移")
    parser.add_argument("--message", default="Initial migration", help="迁移消息")
    parser.add_argument("--upgrade", action="store_true", help="升级数据库")
    
    args = parser.parse_args()
    
    if args.create_migration:
        create_migration(args.message)
    elif args.upgrade:
        upgrade_database()
    else:
        init_database()
